import os
import random
import shutil
import json
import inspect
from datetime import datetime
import msvcrt ,sys ,time
from ultralytics import YOLO
from trans_config import MS_CSV , MS_AVG , MS_END , MS_NML 
from validate import validate_model,show_all_val

# 默认训练配置，可在调用train_model时覆盖
DEFAULT_INIT_PARAMS = {
    'epochs':       110,
    'patience':     40,
    'batch':        0.98,
    'imgsz':        768,
    'device':       0,
    'workers':      2,  

    'augment':      True,
    'optimizer':    'auto', 
    'multi_scale':  False,
    'cos_lr':       False,
    'close_mosaic': 10, 
    'lr0':          0.01,
    'lrf':          0.01,

    'weight_decay': 0.0005,
    'dropout':      0.0,
    'label_smoothing':0.0,
    'warmup_epochs':3,
    'box':          7.5,
    'cls':          0.5,
    'dfl':          1.5,

    'hsv_h':        0.015,  
    'hsv_s':        0.7,  
    'hsv_v':        0.4,  
    'degrees':      0.0,  
    'translate':    0.1,
    'scale':        0.5,  
    'shear':        0.0,  
    'perspective':  0.0,
    'flipud':       0.0,  
    'fliplr':       0.5,  
    'mosaic':       1.0, 
    'mixup':        0.0,  
    'cutmix':       0.0,  
    'copy_paste':   0.0, 
    'erasing':      0.4, 

    'val':          True,
    'plots':        True,
    'amp':          True,
    'cache':        False,
    'verbose':      True,
}

def interpolate_configs(config1, config2, alpha):
    """
    线性插值两个配置字典
    alpha: 插值系数，0->config1, 1->config2
    对于'batch'参数，直接使用config2中的值（如果存在），否则使用config1中的值
    """
    interpolated = {} 
    # 处理两个字典共有的参数
    all_keys = set(config1.keys()) | set(config2.keys())
    for key in all_keys:
        if key == 'batch':
            # 对于'batch'参数，优先使用config2中的值
            interpolated[key] = config2[key] if key in config2 else config1[key]
        elif key in config1 and key in config2 :
            val1 = config1[key]
            val2 = config2[key]
            if isinstance(val1, (int, float)) and isinstance(val2, (int, float)):
                interpolated[key] = val1 + (val2 - val1) * alpha
            else:
                interpolated[key] = val1  # 非数值类型直接使用config1的值
        elif key in config2:
            # 如果参数只在config2中有，直接使用config2的值
            interpolated[key] = config2[key]
        else:
            # 如果参数只在config1中有，直接使用config1的值
            interpolated[key] = config1[key]
    return interpolated

def create_on_train_epoch_start(init_params):
    """
    创建训练轮次开始时的回调函数
    
    """
    def on_train_epoch_start(trainer):
        allocation=[0.1,
                    0.2,
                    0.5,
                    0.2]
        allocated_epochs=[0,0,0,0]
        for index,i in enumerate(allocation):
            allocated_epochs[index]=int(init_params['epochs']*i)
        t1=allocated_epochs[0]
        t2=allocated_epochs[1]
        t3=allocated_epochs[2]
        t4=allocated_epochs[3]
        print(f'分配的训练轮数：{t1,t2,t3,t4}')

        epoch = trainer.epoch
        if epoch <=t1:
            for key, value in MS_CSV.items():
                setattr(trainer, key, value)
            print('='*50)
            print('[回调]正在使用保守配置')
            print('='*50)
        elif t1<epoch<=t1+t2 :
            if t1 < epoch <= t1+10:
                alpha = (epoch - t1) / 10  # 0->t1+1, 1->t1+10   
                config = interpolate_configs(MS_CSV, MS_NML, alpha)
                for key, value in config.items():
                    setattr(trainer, key, value)
                print('='*50)
                print(f'[回调]配置从保守配置过渡到正常配置 (epoch {epoch}, 过渡进度: {alpha:.2f})')
                print('='*50)
            else:
                for key, value in MS_NML.items():
                    setattr(trainer, key, value)
                
                print('='*50)
                print('[回调]配置过渡完成，正在使用正常配置')
                print('='*50)
        elif t1+t2<epoch<=t1+t2+t3 :
            if t1+t2 < epoch <= t1+t2+20:
                alpha = (epoch - t1-t2) / 20  # 0->t1+t2+1, 1->t1+t2+10    
                config = interpolate_configs(MS_NML, MS_AVG, alpha)
                for key, value in config.items():
                    setattr(trainer, key, value)
                print('='*50)
                print(f'[回调]配置从正常配置过渡到激进配置 (epoch {epoch}, 过渡进度: {alpha:.2f})')
                print('='*50)
            else:
                for key, value in MS_AVG.items():
                    setattr(trainer, key, value)
                
                print('='*50)
                print('[回调]配置过渡完成，正在使用激进配置')
                print('='*50)
        elif t1+t2+t3<epoch:
            if t1+t2+t3 < epoch <=t1+t2+t3+20:
                alpha = (epoch - t1-t2-t3) / 20  # 0->t1+t2+t3+1, 1->t1+t2+t3+10    
                config = interpolate_configs(MS_AVG, MS_END, alpha)
                for key, value in config.items():
                    setattr(trainer, key, value)
                print('='*50)
                print(f'[回调]配置从激进配置过渡到收敛配置 (epoch {epoch}, 过渡进度: {alpha:.2f})')
                print('='*50)
            else:
                for key, value in MS_END.items():
                    setattr(trainer, key, value)
                
                print('='*50)
                print('[回调]配置过渡完成，正在使用收敛配置')
                print('='*50)
    return on_train_epoch_start

def progressive_ajust(trainer,hyppara,epoch_range:tuple,value_range:tuple,current_epoch:int):
    '''
        渐进式参数变化
        提供变化初值末值和变化起点终点
    '''
    total_epochs=CUS_INIT_PARAMS['epochs']
    
    if len(epoch_range)==2:
        # 单调变化
        if any(type(i)!=int for i in epoch_range):
            # 以百分比形式设置轮数
            start_epoch,end_epoch=epoch_range

            start_epoch=int(total_epochs*start_epoch)
            end_epoch=int(total_epochs*end_epoch)
        else:
            start_epoch,end_epoch=epoch_range
            start_value,end_value=value_range
        if current_epoch<start_epoch or current_epoch>end_epoch:
            # 不在调整轮数内
            pass
        else:
            contemporary_value=start_value+(end_value-start_value)*(current_epoch-start_epoch)/(end_epoch-start_epoch)
            setattr(trainer,hyppara,contemporary_value)
            print(f'[回调]{hyppara}更新为: {contemporary_value:.4f}')
    elif len(epoch_range)==3:
        # 先单调变化到极值点，再单调变化到目标值
        if any(type(i)!=int for i in epoch_range):
            # 以百分比形式设置轮数
            start_epoch,extreme_epoch,end_epoch=epoch_range

            start_epoch=int(total_epochs*start_epoch)
            extreme_epoch=int(total_epochs*extreme_epoch)
            end_epoch=int(total_epochs*end_epoch)
        else:
            start_epoch,extreme_epoch,end_epoch=epoch_range
            start_value,extreme_value,end_value=value_range
        if current_epoch<start_epoch or current_epoch>end_epoch:
            # 不在调整轮数内
            pass
        elif current_epoch<=extreme_epoch:
            contemporary_value = start_value + (extreme_value-start_value)*(current_epoch-start_epoch)/(extreme_epoch-start_epoch)
            setattr(trainer,hyppara,contemporary_value)
            print(f'[回调]{hyppara}增加为: {contemporary_value:.4f}')
        else:
            contemporary_value = extreme_value + (end_value-extreme_value)*(current_epoch-extreme_epoch)/(end_epoch-extreme_epoch)
            setattr(trainer,hyppara,contemporary_value)
            print(f'[回调]{hyppara}减少为: {contemporary_value:.4f}')
    else:
        print('函数图形错误')

def progressive_HYP_regulation(trainer):
    '''
        集成渐进式参数变化的回调
    '''
    epoch = trainer.epoch
    progressive_ajust(trainer,'mosaic',     (70,140),       (1,0.3),epoch)
    progressive_ajust(trainer,'degrees',    (0.1,0.6,1.0),  (5,8,4),epoch)
    progressive_ajust(trainer,'translate',  (0.1,0.6,1.0),  (0.1,0.15,0.05),epoch)
    progressive_ajust(trainer,'scale',      (0.1,0.6,1.0),  (0.5,0.6,0.4),epoch)
    progressive_ajust(trainer,'mixup',      (0.1,0.6,1.0),  (0.1,0.15,0.05),epoch)
    progressive_ajust(trainer,'cutmix',     (0.1,0.6,1.0),  (0.1,0.15,0.05),epoch)
    progressive_ajust(trainer,'copy_paste', (0.1,0.6,1.0),  (0.1,0.15,0.05),epoch)
    progressive_ajust(trainer,'erasing',    (0.1,0.6,1.0),  (0.4,0.5,0.35),epoch)


def create_on_train_start():
    """
    创建训练开始时的回调函数
    返回:
        回调函数
    """
    def on_train_start(trainer):
        for key, value in MS_CSV.items():
            setattr(trainer, key, value)
        print('='*50)
        print('[回调]配置更新完成')
        print('='*50)
    return on_train_start

def config_save(config_path, ms_configs):
    """
    将训练配置保存为YAML文件（仅包含TRAINING_PARAMS，支持注释）
    
    参数:
    config_path: 保存配置的目标路径
    ms_configs: 包含配置字典的字典，必须包含'TRAINING_PARAMS'键
                可选包含'DEFAULT_PARAMS'键用于生成修改注释
    """
    try:
        # 导入pyyaml库用于YAML处理
        import yaml
        
        # 确保目标目录存在
        os.makedirs(os.path.dirname(config_path), exist_ok=True)
        
        # 检查文件扩展名，如果不是.yaml则添加或修改
        if not config_path.endswith('.yaml') and not config_path.endswith('.yml'):
            new_path = os.path.splitext(config_path)[0] + '.yaml'
            print(f"为支持注释功能，将保存为YAML文件: {new_path}")
            config_path = new_path
        
        # 准备YAML数据，包含注释头部
        yaml_content = []
        yaml_content.append("# 训练配置参数")
        yaml_content.append(f"# 保存时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        yaml_content.append("# 此文件使用YAML格式保存，支持注释功能")
        yaml_content.append("# 与默认参数不同的值会添加'# 修改前: xxx'注释")
        yaml_content.append("")
        
        # 获取训练参数和默认参数（如果存在）
        training_params = ms_configs.get('TRAINING_PARAMS', {})
        default_params = ms_configs.get('DEFAULT_PARAMS', {})
        
        # 确保TRAINING_PARAMS存在
        if not training_params:
            print("警告: TRAINING_PARAMS为空，将创建空配置文件")
        
        # 手动构建带注释的YAML内容 - 直接保存TRAINING_PARAMS的内容，不带TRAINING_PARAMS外层键
        for key, value in training_params.items():
            # 处理可能无法序列化的值
            try:
                yaml.dump(value)
                value_str = str(value)
                # 为布尔值和None特殊处理
                if isinstance(value, bool):
                    value_str = str(value).lower()
                elif value is None:
                    value_str = 'null'
                elif isinstance(value, (int, float)):
                    # 保持数字类型的原始表示
                    pass
                elif isinstance(value, str):
                    # 如果字符串包含特殊字符，需要加引号
                    if any(c in value for c in ': #{}[]&*!|\'"<>=') or ' ' in value:
                        # 处理可能包含单引号的情况
                        if "'" in value and '"' not in value:
                            value_str = f'"{value}"'
                        else:
                            value_str = f"'{value}'"
            except (TypeError, OverflowError):
                # 将不可序列化的值转换为字符串
                value_str = f"'{str(value)}'"
            
            # 添加注释（如果与默认值不同）
            comment = ""
            if key in default_params and default_params[key] != value:
                comment = f"  # 修改前: {default_params[key]}"
            
            yaml_content.append(f"{key}: {value_str}{comment}")
        
        # 保存为YAML文件
        with open(config_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(yaml_content))
        
        print(f"配置已保存到YAML文件: {config_path}")
        return True
    except ImportError:
        # 如果pyyaml未安装，提供友好提示
        print("错误：需要安装pyyaml库来保存YAML格式配置。请运行 'pip install pyyaml'")
        return False
    except Exception as e:
        print(f"保存配置时出错: {str(e)}")
        return False

def convert_json_to_yolo_txt(json_file_path, txt_file_path, class_mapping):
    """
    将LabelMe JSON文件转换为YOLO格式的TXT文件
    
    参数:
    json_file_path: JSON文件路径
    txt_file_path: 输出的TXT文件路径
    class_mapping: 类别名称到ID的映射字典
    """
    try:
        # 读取JSON文件
        with open(json_file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # 获取图像尺寸
        img_width = data['imageWidth']
        img_height = data['imageHeight']
        
        # 准备YOLO格式内容
        yolo_lines = []
        
        # 处理每个标注形状
        for shape in data['shapes']:
            label = shape['label']
            points = shape['points']
            shape_type = shape.get('shape_type', 'polygon')
            
            # 获取类别ID
            if label not in class_mapping:
                # 如果类别不在映射中，跳过或处理错误
                print(f"警告: 在文件 {json_file_path} 中发现未知类别 '{label}'，已跳过")
                continue
                
            class_id = class_mapping[label]
            
            # 根据形状类型处理坐标
            if shape_type in ['polygon', 'rectangle']:
                # 多边形和矩形：计算外接矩形
                # 对于矩形，LabelMe可能存储为2个对角点
                # 使用所有点来计算最小外接矩形以确保准确性
                xs = [p[0] for p in points]
                ys = [p[1] for p in points]
                
                x_min = min(xs)
                x_max = max(xs)
                y_min = min(ys)
                y_max = max(ys)
            else:
                print(f"警告: 在文件 {json_file_path} 中发现不支持的形状类型 '{shape_type}'，已跳过")
                continue
            
            # 计算中心点、宽度和高度，并归一化
            x_center = (x_min + x_max) / 2 / img_width
            y_center = (y_min + y_max) / 2 / img_height
            width = (x_max - x_min) / img_width
            height = (y_max - y_min) / img_height
            
            # 添加到YOLO格式内容
            yolo_lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
            
        
        # 写入TXT文件
        with open(txt_file_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(yolo_lines))
            
        return True
        
    except Exception as e:
        print(f"错误: 处理文件 {json_file_path} 时出错: {str(e)}")
        return False

def create_dataset_config(target_dir, train_count, val_count, classes):
    """创建YOLO数据集配置文件"""
    config_path = os.path.join(target_dir, "dataset.yaml")
    
    # 生成类别名称部分 - 注意缩进格式必须正确
    names_content = ""
    for idx, cls in enumerate(classes):
        names_content += f"  {idx}: {cls}\n"
    
    # 修复YAML格式错误：1)去除所有额外缩进 2)将nc参数移到正确位置
    config_content = f"""# YOLO 数据集配置文件
# 由自动整理脚本生成于 {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

# 数据集路径
path: {os.path.abspath(target_dir)}
train: images/train  # 训练图像路径 (相对path)
val: images/val      # 验证图像路径 (相对path)

# 类别信息
names:
{names_content}nc: {len(classes)}     # 类别数量

# 数据集统计
# 训练样本数: {train_count}
# 验证样本数: {val_count}
# 总样本数: {train_count + val_count}
"""
    
    with open(config_path, 'w', encoding='utf-8') as f:
        f.write(config_content)
    
    print(f"\n数据集配置文件已创建: {config_path}")

def organize_dataset(source_dir, target_dir, val_ratio=0.2):
    """
    整理数据集并划分训练集和验证集
    
    参数:
    source_dir: 包含混合图像和标注文件的源目录
    target_dir: 整理后的目标目录
    val_ratio: 验证集比例 (默认0.2即20%)
    """
    # 确保目标目录存在
    os.makedirs(target_dir, exist_ok=True)
    
    # 创建标准YOLO目录结构
    dirs = {
        'images_train': os.path.join(target_dir, 'images', 'train'),
        'images_val': os.path.join(target_dir, 'images', 'val'),
        'labels_train': os.path.join(target_dir, 'labels', 'train'),
        'labels_val': os.path.join(target_dir, 'labels', 'val'),
    }
    
    for d in dirs.values():
        os.makedirs(d, exist_ok=True)
    
    # 获取所有图像文件
    image_exts = ['.jpg', '.jpeg', '.png', '.bmp']
    image_files = [f for f in os.listdir(source_dir) 
                  if os.path.splitext(f)[1].lower() in image_exts]
    
    # 分离有标注和无标注的图像
    paired_files = []  # 正样本（有标注）
    negative_samples = []  # 负样本（无标注）
    
    for img_file in image_files:
        base_name = os.path.splitext(img_file)[0]
        json_file = base_name + '.json'
        if os.path.exists(os.path.join(source_dir, json_file)):
            paired_files.append((img_file, json_file))
        else:
            negative_samples.append(img_file)
    
    print(f"找到 {len(paired_files)} 个带标注的图像和 {len(negative_samples)} 个无标注的图像")
    
    # 如果没有找到任何带标注的图像，提示错误
    if not paired_files:
        print("错误: 未找到任何JSON标注文件!")
        return
    
    # 从所有JSON文件中提取所有类别
    all_classes = set()
    for img_file, json_file in paired_files:
        json_path = os.path.join(source_dir, json_file)
        try:
            with open(json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                for shape in data['shapes']:
                    all_classes.add(shape['label'])
        except Exception as e:
            print(f"警告: 读取文件 {json_path} 时出错: {str(e)}")
    
    # 创建类别映射（按字母顺序排序）
    sorted_classes = sorted(list(all_classes))
    class_mapping = {cls: idx for idx, cls in enumerate(sorted_classes)}
    
    print(f"找到 {len(all_classes)} 个类别: {sorted_classes}")
    print(f"类别映射: {class_mapping}")
    
    # 随机打乱文件对
    random.shuffle(paired_files)
    
    # 计算验证集数量
    total_count = len(paired_files)
    val_count = int(total_count * val_ratio)
    train_count = total_count - val_count
    
    print(f"找到 {total_count} 个图像-标注对")
    print(f"训练集: {train_count} 个样本")
    print(f"验证集: {val_count} 个样本")
    
    # 处理并复制带标注的图像（正样本）到相应目录
    for i, (img_file, json_file) in enumerate(paired_files):
        base_name = os.path.splitext(img_file)[0]
        txt_file = base_name + '.txt'
        
        # 转换JSON到YOLO TXT格式
        json_path = os.path.join(source_dir, json_file)
        
        if i < train_count:
            # 训练集
            txt_path = os.path.join(dirs['labels_train'], txt_file)
            img_dest = os.path.join(dirs['images_train'], img_file)
        else:
            # 验证集
            txt_path = os.path.join(dirs['labels_val'], txt_file)
            img_dest = os.path.join(dirs['images_val'], img_file)
        
        # 转换JSON到YOLO格式
        if convert_json_to_yolo_txt(json_path, txt_path, class_mapping):
            # 复制图像文件
            shutil.copy(os.path.join(source_dir, img_file), img_dest)
            print(f"处理正样本文件 {i+1}/{total_count}: {img_file}")
    
    # 处理无标注的图像（负样本）
    if negative_samples:
        # 随机打乱负样本
        random.shuffle(negative_samples)
        
        # 计算负样本在训练集和验证集的分配数量
        negative_val_count = int(len(negative_samples) * val_ratio)
        negative_train_count = len(negative_samples) - negative_val_count
        
        print(f"负样本分配: 训练集 {negative_train_count} 个, 验证集 {negative_val_count} 个")
        
        # 复制负样本到相应目录
        for i, img_file in enumerate(negative_samples):
            if i < negative_train_count:
                # 训练集
                img_dest = os.path.join(dirs['images_train'], img_file)
                # 为负样本创建空的标签文件，YOLO要求图像和标签文件一一对应
                txt_file = os.path.splitext(img_file)[0] + '.txt'
                txt_dest = os.path.join(dirs['labels_train'], txt_file)
                open(txt_dest, 'a').close()  # 创建空文件
            else:
                # 验证集
                img_dest = os.path.join(dirs['images_val'], img_file)
                # 为负样本创建空的标签文件
                txt_file = os.path.splitext(img_file)[0] + '.txt'
                txt_dest = os.path.join(dirs['labels_val'], txt_file)
                open(txt_dest, 'a').close()  # 创建空文件
            
            # 复制图像文件
            shutil.copy(os.path.join(source_dir, img_file), img_dest)
            print(f"处理负样本文件 {i+1}/{len(negative_samples)}: {img_file}")
    
    # 更新最终的训练集和验证集数量，包括正负样本
    final_train_count = train_count + (len(negative_samples) - int(len(negative_samples) * val_ratio))
    final_val_count = val_count + int(len(negative_samples) * val_ratio)
    
    # 创建数据集配置文件
    create_dataset_config(target_dir, final_train_count, final_val_count, sorted_classes)
    
    print(f"\n数据集整理完成！保存在: {target_dir}")
    print("请检查目录结构并开始训练！")
    print(f"最终数据集组成: 训练集 {final_train_count} 个样本(正样本{train_count}个,负样本{negative_train_count}个), 验证集 {final_val_count} 个样本(正样本{val_count}个,负样本{negative_val_count}个)")

def train_model(source_dir="captured_photos", model_path="yolo11s.pt", 
    custom_params=None, experiment_name=None,send_email=False,automatic=False):
    """
    训练YOLO模型的核心函数
    
    参数:
        source_dir: 包含混合图像和标注文件的源目录
        model_path: 预训练模型文件路径
        custom_params: 自定义训练参数字典，用于覆盖默认参数
        experiment_name: 实验名称，如果为None则自动生成
    
    返回:
        验证结果
    """
    # 创建带时间戳的子目录，避免覆盖之前的数据集
    timestamp = datetime.now().strftime("%m%d_%H%M")  # 简化为月日_时分格式
    target_dir = os.path.join("yolo_dataset", timestamp)  # 在yolo_dataset下创建带时间戳的子目录
    
    # 执行数据集整理
    organize_dataset(source_dir, target_dir)

    # 检查本地模型是否存在
    full_model_path = model_path if os.path.isabs(model_path) else f'E:/Py/visual/{model_path}'
    if os.path.exists(full_model_path):
        print(f"本地模型文件: {full_model_path}")
        model = YOLO(full_model_path, task='detect', verbose=True)
    else:
        print(f"本地模型文件不存在: {full_model_path}")
        return None
    
    # 确定实验名称
    if experiment_name is None:
        experiment_name = timestamp
    else:
        experiment_name = f"{experiment_name}_{timestamp}"

    print(f'输出为: {experiment_name}')
    # 当custom_params不为空时，对比并打印变化的参数
    if custom_params:
        print('='*50)
        print('自定义参数对比:')
        for key, new_value in custom_params.items():
            if key in DEFAULT_INIT_PARAMS and DEFAULT_INIT_PARAMS[key] != new_value:
                old_value = DEFAULT_INIT_PARAMS[key]
                print(f'{key} : {old_value} -> {new_value}')
        print('='*50)
    if not automatic:
        print('按回车开始训练，按 ESC 退出程序')
        while True:
            if msvcrt.kbhit():  # 检测是否有按键
                key = msvcrt.getch()
                if key == b'\r':  # 回车键
                    break
                elif key == b'\x1b':  # ESC 键
                    sys.exit(0)
    else:
        pass

    # 获取数据集配置路径
    dataset_config = os.path.abspath(os.path.join(target_dir, 'dataset.yaml'))
    
    # 创建训练参数 - 基于默认参数并覆盖自定义参数
    init_params = DEFAULT_INIT_PARAMS.copy()
    # 设置必要的参数
    init_params['data'] = dataset_config
    init_params['name'] = experiment_name
    # 应用自定义参数
    if custom_params:
        init_params.update(custom_params)
    
    # 创建回调函数
    model.add_callback('on_train_epoch_start', progressive_HYP_regulation)
    # model.add_callback('on_train_start', create_on_train_start)
    
    # 开始训练
    t = time.time()
    results = model.train(**init_params)
    training_time = int((time.time() - t) / 60)
    
    # 验证模型
    best_model_path = f'E:/Py/visual/runs/detect/{experiment_name}/weights/best.pt'
    rs = validate_model(best_model_path, t=training_time)
    
    
    
    # 删除临时数据集目录，节省空间
    try:
        if os.path.exists(target_dir):
            shutil.rmtree(target_dir)
            print(f"已删除临时数据集目录: {target_dir}")
        else:
            print(f"临时数据集目录不存在: {target_dir}")
    except Exception as e:
        print(f"删除临时数据集目录时出错: {str(e)}")
    
    show_all_val('E:/Py/visual/runs/detect',block=send_email)
    print(f'训练时长: {training_time}分钟')
    
    # 训练完成后保存配置到YAML文件
    if custom_params is not None:
        # 构建ms_configs格式，包含自定义参数
        ms_configs = {
            'TRAINING_PARAMS': custom_params,
            'DEFAULT_PARAMS': DEFAULT_INIT_PARAMS  # 同时保存默认参数用于比较
        }
        
        # 创建保存路径，使用实验名称
        config_path = f'E:/Py/visual/runs/detect/{experiment_name}/training_config.yaml'
        
        # 调用配置保存函数
        config_save(config_path, ms_configs)
    
    return rs ,experiment_name

def main(custom_params=None,send_email=False,exp_name=None):
    """
    主函数入口
    提供交互式界面，调用核心训练函数
    """
    source_dir = "captured_photos"  # 默认源目录
    pre_model = 'yolo11s.pt'  # 默认预训练模型
    automatic=False        
    # 获取实验名称
    if exp_name is None:
        n = input('输入本次实验名称，按下回车跳过以使用默认名称')
        experiment_name = None if n.strip() == '' else n.strip()
    else:
        automatic=True
        experiment_name = exp_name
    # 调用核心训练函数
    rs ,exp_name= train_model(source_dir=source_dir,
                    model_path=pre_model,
                    experiment_name=experiment_name,
                    custom_params=custom_params,
                    send_email=not send_email,
                    automatic=automatic
                    )
    if send_email:
        from tran_fin_inform import send_qq_email, generate_training_report_html
        import datetime
        import base64
        import os
        
        # 读取图片二进制数据
        results_path = f'E:/Py/visual/runs/detect/{exp_name}/results.png'
        img_base64 = ''
        image_data = None
        
        if os.path.exists(results_path):
            print(f"图片文件存在: {results_path}")
            try:
                file_size = os.path.getsize(results_path)
                print(f"图片文件大小: {file_size} 字节")
                
                with open(results_path, 'rb') as img_file:
                    image_data = img_file.read()
                    print(f"成功读取图片数据，长度: {len(image_data)} 字节")
                    # 只在需要时进行base64编码
                    if image_data:
                        img_base64 = base64.b64encode(image_data).decode('utf-8')
            except Exception as e:
                print(f"读取图片失败: {e}")
        else:
            # 尝试查找其他可能的结果图路径
            import glob
            possible_images = glob.glob(f'E:/Py/visual/runs/detect/{exp_name}/*.png')
            print(f"在实验目录中找到的PNG文件: {possible_images}")
        
        # 使用外部函数生成HTML报告内容
        html_message = generate_training_report_html(
            results_dict=rs,  # 训练结果字典
            experiment_name=exp_name,  # 实验名称
            img_base64=img_base64  # 传递base64仅用于函数内部逻辑判断
        )
        
        # 发送邮件
        time_stamp = datetime.datetime.now().strftime("%m%d_%H%M")  # 简化为月日_时分格式
        send_qq_email(
            subject=f'训练在{time_stamp}完成!', 
            message=html_message,
            to_addr='2264261565@qq.com',
            is_html=True,
            image_data=image_data,  # 传递原始二进制图片数据
            image_filename="训练结果图表.png"
        )

    return rs

if __name__ == "__main__":
    global CUS_INIT_PARAMS
    CUS_INIT_PARAMS = {
    'epochs':       150,
    'patience':     40,
    'batch':        0.98,
    'imgsz':        800,
    'device':       0,
    'workers':      2,  

    'augment':      True,
    'optimizer':    'AdamW', 
    'multi_scale':  False,
    'cos_lr':       True,
    'close_mosaic': 00, 
    'lr0':          8e-4,
    'lrf':          0.07,

    'weight_decay': 0.001,
    'dropout':      0.005,
    'label_smoothing':0.0,

    'warmup_epochs':3,
    'box':          7.5,
    'cls':          0.5,
    'dfl':          1.5,

    'hsv_h':        0.015,  
    'hsv_s':        0.7,  
    'hsv_v':        0.4,  
    'degrees':      5.0,  
    'translate':    0.1,
    'scale':        0.5,  
    'shear':        0.1,  
    'perspective':  0.0,
    'flipud':       0.1,  
    'fliplr':       0.5,  
    'mosaic':       1.0, 
    'mixup':        0.1,  
    'cutmix':       0.1,  
    'copy_paste':   0.1, 
    'erasing':      0.4, 

    'val':          True,
    'plots':        True,
    'amp':          True,
    'cache':        False,
    'verbose':      True,
}

    rs=main(custom_params=CUS_INIT_PARAMS,send_email=True,exp_name='v3_PGS')
    
    
    # 三个训练任务完成后关闭电脑
    # import os
    # print("所有训练任务已完成，准备关闭电脑...")
    # os.system("shutdown /s /t 0")
    

